package http.secure.common.filter;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import http.secure.common.anno.DisguiseAPI;
import http.secure.common.anno.ResponseEnc;
import http.secure.common.constant.DisguiseEnum;
import http.secure.common.constant.SecureConstant;
import http.secure.common.constant.SecureProperties;
import http.secure.common.exception.InvalidResponseException;
import http.secure.common.keys.KeySwitcher;
import http.secure.common.keys.SecureKeyStore;
import http.secure.common.util.Assert;
import http.secure.common.util.SMHelper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.MethodParameter;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.core.io.Resource;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

/**
 * HttpServletResponse cryptographic processing
 */
@Slf4j
@ControllerAdvice
@RequiredArgsConstructor
public class EncryptResponseBody implements ResponseBodyAdvice<Object> {
    private final SecureKeyStore keyStore;
    private final ObjectMapper objectMapper;
    private final KeySwitcher keySwitcher;
    private final SecureProperties properties;

    @Override
    public boolean supports(MethodParameter parameter, Class<? extends HttpMessageConverter<?>> convert) {
        return true;
    }

    public Object beforeBodyWrite(Object body, MethodParameter parameter, MediaType mediaType, Class<? extends HttpMessageConverter<?>> converterType, ServerHttpRequest request, ServerHttpResponse response) {
        // 检查 DisguiseAPI 注解
        DisguiseAPI disguiseAPI = parameter.getMethodAnnotation(DisguiseAPI.class);
        if (disguiseAPI != null) {
            response.getHeaders().add("X-EXCHANGE-API", properties.getApi().getKeyExchangePath());
            response.getHeaders().add("X-ROTATION-API", properties.getApi().getKeyRotationPath());
            return injectDisguise(body, disguiseAPI.value());
        }
        // 检查 ResponseEnc 注解
        ResponseEnc responseEnc = parameter.getMethodAnnotation(ResponseEnc.class);
        if (responseEnc == null) {
            responseEnc = parameter.getDeclaringClass().getAnnotation(ResponseEnc.class);
        }
        if (responseEnc != null) {
            return responseEnc(body, mediaType, response);
        }
        return body;
    }

    private Object injectDisguise(Object body, DisguiseEnum disguise) {
        // 仅支持org.springframework.core.io.Resource类型
        Assert.isTrue(body instanceof Resource, "Body must be of type Resource", InvalidResponseException::new);
        // 读取二进制数据
        try {
            byte[] bytes = ((Resource) body).getContentAsByteArray();
            byte[] disguiseData = keySwitcher.disguisePubkey(disguise, bytes);
            return new ByteArrayResource(disguiseData);
        } catch (Exception e) {
            throw new InvalidResponseException(e.getMessage());
        }
    }

    // 加密响应
    private Object responseEnc(Object body, MediaType mediaType, ServerHttpResponse response) {
        // 仅支持application/json媒体类型
        boolean isJson = MediaType.APPLICATION_JSON.isCompatibleWith(mediaType);
        boolean isText = MediaType.TEXT_PLAIN.isCompatibleWith(mediaType);
        Assert.isTrue(isJson || isText, "Response type not support encrypt", InvalidResponseException::new);
        // 轮换前检查
        beforeRotationCheck(response);
        // 获取 SM4 密钥
        String key = keyStore.get(SecureConstant.SM4_KEY);
        String iv = keyStore.get(SecureConstant.BROWSER_IV);
        Assert.notNull(key, "Key is null", InvalidResponseException::new);
        Assert.notNull(iv, "IV is null", InvalidResponseException::new);
        if (isJson) {
            // 序列化JSON
            try {
                body = objectMapper.writeValueAsString(body);
            } catch (JsonProcessingException e) {
                throw new InvalidResponseException("JSON serialization failed: " + e.getMessage());
            }
        }
        // 加密响应
        return SMHelper.sm4Encrypt(String.valueOf(body), key, iv);
    }

    /**
     * 新旧密钥轮换
     */
    private void beforeRotationCheck(ServerHttpResponse response) {
        // 密钥轮换时间小于等于0表示永不过期（不推荐）
        long rotationTime = properties.getKeyRotationTime();
        if (rotationTime <= 0) {
            return;
        }
        Long expireTime = keyStore.getExpireTime(SecureConstant.SM4_KEY);
        // 判断旧密钥是否过期
        if (expireTime == null || rotationTime * 1000 > expireTime) {
            String cpk = keyStore.get(SecureConstant.CLIENT_PUB_KEY);
            // 获取新密钥
            String newKey = keySwitcher.rotationKeys(cpk);
            // 通过响应头推送
            response.getHeaders().add("X-EXCHANGE-KEY", newKey);
        }
    }
}
